from chromadb import EmbeddingFunction
from chromadb import PersistentClient

class MyEmbeddingFunction(EmbeddingFunction):
    def __call__(self, input: list):
        from sentence_transformers import SentenceTransformer
        model = SentenceTransformer('BAAI/bge-m3')
        return [model.encode(text).tolist() for text in input]

embedding_function = MyEmbeddingFunction()
chroma_client = PersistentClient(path="./chroma_db")
collections = {}

import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def getCollection(id):
    collection_name = "doc_" + str(id)
    if collection_name not in collections:
        logging.info(f'正在创建集合: {collection_name}')
        collections[collection_name] = chroma_client.get_or_create_collection(name=collection_name, embedding_function=embedding_function)
    else:
        logging.info(f'正在使用已有的集合: {collection_name}')
    return collections[collection_name]
    